from diffusers import StableDiffusionPipeline
import torch
import os

def download_model():
    """下载 Stable Diffusion v1.5 模型"""
    print("开始下载 Stable Diffusion v1.5 模型...")
    print("模型大小约为 4GB，请确保有足够的磁盘空间")
    
    # 设置模型ID
    model_id = "runwayml/stable-diffusion-v1-5"
    
    try:
        # 下载模型
        pipe = StableDiffusionPipeline.from_pretrained(
            model_id,
            torch_dtype=torch.float16,
            use_safetensors=True
        )
        
        # 保存到本地
        save_path = "./stable-diffusion-v1-5"
        pipe.save_pretrained(save_path)
        
        print(f"模型下载完成！保存在: {os.path.abspath(save_path)}")
        print("你现在可以在 stable_diffusion_demo.py 中使用这个模型了")
        
    except Exception as e:
        print(f"下载过程中发生错误: {str(e)}")
        print("请检查网络连接和磁盘空间")

if __name__ == "__main__":
    download_model() 